#ifndef UnstructMG_CSR_MAT_OP
#define UnstructMG_CSR_MAT_OP

#include "seq_mv.hpp"
#include "scal_kernels.hpp"
#include "qsort.hpp"
#include "binsearch.hpp"

template<typename idx_t,    typename data_t1, typename calc_t1,
                            typename data_t2, typename calc_t2, int dof>
void CSRMatrixTranspose(CSRMatrix<idx_t, data_t1, calc_t1, dof> & src,
                        CSRMatrix<idx_t, data_t2, calc_t2, dof> & dst)
{
    assert(src.nnz == src.row_ptr[src.nrows]);
    if (src.nrows > 0 && src.ncols == 0) {// 确定列数
        idx_t max_col = -1;
        for (idx_t i = 0; i < src.nrows; i ++) {
            for (idx_t p = src.row_ptr[i]; p < src.row_ptr[i + 1]; p ++) {
                if (src.col_idx[p] > max_col) {
                    max_col = src.col_idx[p];
                }
            }
        }
        src.ncols = max_col + 1;
    }

    dst.nrows = src.ncols;
    dst.ncols = src.nrows;
    dst.nnz   = src.nnz  ;
    const idx_t num_cols_A = src.ncols;

    if (dst.nrows == 0) return ;

    dst.alloc_mem(false);
    idx_t * bucket = new idx_t [(num_cols_A + 1) * omp_get_max_threads()];

    // idx = idx2*dim1 + idx1  ->  ret = idx1*dim2 + idx2 = (idx%dim1)*dim2 + idx/dim1
    auto transpose_idx = [](const idx_t idx, const idx_t dim1, const idx_t dim2) {
        return idx % dim1 * dim2 + idx / dim1;
    };
    constexpr long long e_size = dof*dof;
    
    #pragma omp parallel
    {
        int tid = omp_get_thread_num();
        int nt  = omp_get_num_threads();
        int ibeg, iend;
        thread_load_balance(src.nrows, nt, tid, ibeg, iend);

        // 初始化计数
        #pragma omp for schedule(static)
        for (idx_t i = 0; i < (num_cols_A + 1) * nt; i++)
            bucket[i] = 0;

        // Count the number of entries that will go into each bucket
        // bucket is used as HYPRE_Int[num_threads][num_colsA] 2D array
        for (idx_t p = src.row_ptr[ibeg]; p < src.row_ptr[iend]; p ++) {
            bucket[tid * num_cols_A + src.col_idx[p]] ++;
        }

        // Parallel prefix sum of bucket with length num_colsA * num_threads
        // accessed as if it is transposed as HYPRE_Int[num_colsA][num_threads]
        #pragma omp barrier
        for (idx_t i = tid * num_cols_A + 1; i < (tid + 1) * num_cols_A; i ++) {
            idx_t transpose_i = transpose_idx(i, nt, num_cols_A);
            idx_t transpose_i_minus_1 = transpose_idx(i - 1, nt, num_cols_A);
            bucket[transpose_i] += bucket[transpose_i_minus_1];
        }
        #pragma omp barrier
        #pragma omp master
        {
            for (int i = 1; i < nt; i ++) {
                idx_t j0 = num_cols_A *  i      - 1;
                idx_t j1 = num_cols_A * (i + 1) - 1;
                idx_t transpose_j0 = transpose_idx(j0, nt, num_cols_A);
                idx_t transpose_j1 = transpose_idx(j1, nt, num_cols_A);
                bucket[transpose_j1] += bucket[transpose_j0];
            }
        }
        #pragma omp barrier
        if (tid > 0) {
            idx_t transpose_i0 = transpose_idx(num_cols_A * tid - 1, nt, num_cols_A);
            idx_t offset = bucket[transpose_i0];
            for (idx_t i = tid * num_cols_A; i < (tid + 1) * num_cols_A - 1; i ++) {
                idx_t transpose_i = transpose_idx(i, nt, num_cols_A);
                bucket[transpose_i] += offset;
            }
        }
        #pragma omp barrier

        // Load the data and column numbers of AT
        for (idx_t i = iend - 1; i >= ibeg; i --) {
            for (idx_t p = src.row_ptr[i + 1] - 1; p >= src.row_ptr[i]; p --) {
                idx_t idx = src.col_idx[p];
                bucket[tid * num_cols_A + idx] --;

                idx_t offset = bucket[tid * num_cols_A + idx];
                dst.col_idx[offset] = i;
                for (int r = 0; r < dof; r++)
                for (int c = 0; c < dof; c++)
                    dst.vals[offset * e_size + c * dof + r] = src.vals[p * e_size + r * dof + c];
            }
        }
        #pragma omp barrier
        
        #pragma omp for schedule(static)
        for (idx_t i = 0; i < num_cols_A + 1; i++) {
            dst.row_ptr[i] = bucket[i];
        }
    }// omp region

    dst.row_ptr[num_cols_A] = dst.nnz;

    delete [] bucket;
}


template<typename idx_t,    typename data_t, typename calc_t, int dof1, int dof2, int dof3>
void CSRMatrixMultiply( const CSRMatrix<idx_t, data_t, calc_t, dof1> & A,
                        const CSRMatrix<idx_t, data_t, calc_t, dof2> & B, 
                        CSRMatrix<idx_t, data_t, calc_t, dof3> & C)
{
    static_assert(dof1 == dof2 || dof1 == 1 || dof2 == 1);
    static_assert(dof3 == (MAX(dof1, dof2)));
    assert(A.ncols == B.nrows);
    const bool allsquare = A.nrows == B.ncols;
    
    C.nrows = A.nrows;
    C.ncols = B.ncols;

    if (A.nnz == 0 || B.nnz == 0) {    
        C.nnz   = 0;
        C.alloc_mem();
        return ;
    }

    // Allocate memory
    int * twspace = new int [omp_get_max_threads()];
    C.row_ptr = new idx_t [A.nrows + 1];

    #pragma omp parallel
    {
        #pragma omp for schedule(static)
        for (idx_t i = 0; i < A.nrows + 1; i++)
            C.row_ptr[i] = 0;

        int tid = omp_get_thread_num();
        int  nt = omp_get_num_threads();
        idx_t ibeg, iend;
        thread_load_balance(A.nrows, nt, tid, ibeg, iend);

        idx_t * B_marker = new idx_t [B.ncols];// 记录第j列是否有贡献非零元
        for (idx_t j = 0; j < B.ncols; j++) // 初始化标记
            B_marker[j] = -1;
        
        // ======================================
        // First Pass: compute sizes of C rows. 
        // ======================================

        idx_t num_nonzeros = 0;
        for (idx_t ic = ibeg; ic < iend; ic++) {
            idx_t iic = -1;
            if (false) {
                // iic = rownnz_A[ic];
                // C_i[iic] = num_nonzeros;
            } else {
                iic = ic;
                C.row_ptr[iic] = num_nonzeros;
                if (allsquare) {
                    B_marker[iic] = iic;
                    num_nonzeros ++;
                }
            }

            for (idx_t pa = A.row_ptr[iic]; pa < A.row_ptr[iic + 1]; pa ++) {
                idx_t ja = A.col_idx[pa];
                for (idx_t pb = B.row_ptr[ja]; pb < B.row_ptr[ja + 1]; pb ++) {
                    idx_t jb = B.col_idx[pb];
                    if (B_marker[jb] != iic) {// 此时marker数组在这一列记录的是上次某一行的
                        B_marker[jb] = iic;
                        num_nonzeros++;
                    }
                }
            }
        }
        twspace[tid] = num_nonzeros;

        #pragma omp barrier

        // Correct C_i - phase 1
        if (tid > 0) {
            idx_t jj = twspace[0];
            for (idx_t i1 = 1; i1 < tid; i1++) {
                jj += twspace[i1];
            }
            for (idx_t i1 = ibeg; i1 < iend; i1++) {
                // idx_t iic = rownnz_A ? rownnz_A[i1] : i1;
                idx_t iic = i1;
                C.row_ptr[iic] += jj;
            }
        } else {// 主线程负责统计整体大小并开辟内存
            C.row_ptr[A.nrows] = 0;
            for (idx_t i1 = 0; i1 < nt; i1++)
            {
                C.row_ptr[A.nrows] += twspace[i1];
            }
            C.nnz = C.row_ptr[A.nrows];
            C.col_idx = new idx_t [C.nnz]; 
            C.vals    = new data_t[C.nnz * dof3 * dof3];
        }

        // Correct C_i - phase 2 
        // if (rownnz_A != nullptr) {
        // }
        // End of First Pass
        #pragma omp barrier

        // ====================================
        // Second pass: Fill in C_data and C_j.
        // ==================================== 

        for (idx_t j = 0; j < B.ncols; j++) {// 重新初始化标记
            B_marker[j] = -1;
        }

        // idx_t counter = rownnz_A ? C_i[rownnz_A[ns]] : C_i[ns];
        idx_t counter = C.row_ptr[ibeg];
        for (idx_t ic = ibeg; ic < iend; ic++) {
            idx_t iic = -1;
            if (false) {
                // iic = rownnz_A[ic];
            } else {
                iic = ic;
                if (allsquare) {
                    B_marker[ic] = counter;
                    C.col_idx[counter] = ic;// C_j[counter] = ic;
                    for (int f = 0; f < dof3 * dof3; f++)// C_data[counter] = 0;
                        C.vals[counter * dof3 * dof3 + f] = 0.0;
                    counter++;
                }
            }

            for (idx_t pa = A.row_ptr[iic]; pa < A.row_ptr[iic + 1]; pa ++) {
                idx_t ja = A.col_idx[pa];
                data_t * a_ptr = A.vals + pa * dof1 * dof1;// a_entry = A_data[ia];
                for (idx_t pb = B.row_ptr[ja]; pb < B.row_ptr[ja + 1]; pb ++) {
                    idx_t jb = B.col_idx[pb];
                    data_t * b_ptr = B.vals + pb * dof2 * dof2;// b_entry = B_data[ib];
                    if (B_marker[jb] < C.row_ptr[iic]) {// 此时marker数组在这一列记录的是上次某一行的
                        B_marker[jb] = counter;
                        C.col_idx[counter] = jb;// 第一次在iic行算到该列，记录列序号
                        data_t * c_ptr = C.vals + counter * dof3 * dof3;// C_data[B_marker[jb]] = a_entry * b_entry;
                        if constexpr (dof1 == 1) {
                            for (int f = 0; f < dof2 * dof2; f++)
                                c_ptr[f] = a_ptr[0] * b_ptr[f];
                        } else if constexpr (dof2 == 1) {
                            for (int f = 0; f < dof1 * dof1; f++)
                                c_ptr[f] = a_ptr[f] * b_ptr[0];
                        } else { static_assert(dof1 == dof2);
                            matmat_mul<data_t, data_t, data_t, dof1>(1.0, a_ptr, b_ptr, c_ptr);
                        }
                        counter++;
                    }
                    else {// 此时marker数组的记录是对应本行计算的
                        data_t * c_ptr = C.vals + B_marker[jb] * dof3 * dof3;// C_data[B_marker[jb]] = a_entry * b_entry;
                        if constexpr (dof1 == 1) {
                            for (int f = 0; f < dof2 * dof2; f++)
                                c_ptr[f] += a_ptr[0] * b_ptr[f];
                        } else if constexpr (dof2 == 1) {
                            for (int f = 0; f < dof1 * dof1; f++)
                                c_ptr[f] += a_ptr[f] * b_ptr[0];
                        } else { static_assert(dof1 == dof2);
                            matmat_mla<data_t, dof1>(1.0, a_ptr, b_ptr, 1.0, c_ptr);
                        }
                        // C_data[B_marker[jb]] += a_entry * b_entry;
                    }
                }
            }
        }
        // End of second pass
        delete [] B_marker;

        // 每行升序排列
        for (idx_t i = ibeg; i < iend; i++) {
            sort_indices_with_vals<idx_t, data_t, dof3>(C.row_ptr[i + 1] - C.row_ptr[i],
                C.col_idx + C.row_ptr[i], C.vals + C.row_ptr[i] * dof3 * dof3);
        }
    }// omp region

    delete [] twspace;
}

// 不保证结果矩阵每一行内列序号升序排列
template<typename idx_t,    typename data_t, typename calc_t, int dof>
void CSRMatrixAdd(  const calc_t alpha, const CSRMatrix<idx_t, data_t, calc_t, dof> & A,
                    const calc_t beta , const CSRMatrix<idx_t, data_t, calc_t, dof> & B,
                    CSRMatrix<idx_t, data_t, calc_t, dof> & C)
{
    assert(A.nrows == B.nrows && A.ncols == B.ncols);
    constexpr int e_size = dof*dof;

    idx_t * twspace = new idx_t [omp_get_max_threads()];
    idx_t nnzrows_C = A.nrows;
    // idx_t * rownnz_C = nullptr;

    C.nrows = A.nrows;
    C.ncols = A.ncols;
    C.row_ptr = new idx_t [C.nrows + 1];

    #pragma omp parallel
    {
        #pragma omp for schedule(static)
        for (idx_t i = 0; i <  C.nrows + 1; i++)
            C.row_ptr[i] = 0;

        int tid = omp_get_thread_num();
        int  nt = omp_get_num_threads();
        idx_t ibeg = -1, iend = -1;
        idx_t * marker = nullptr;

        thread_load_balance(nnzrows_C, nt, tid, ibeg, iend);
        
        marker = new idx_t [C.ncols];
        for (idx_t j = 0; j < C.ncols; j++) marker[j] = -1;// 初始化标记
        
        idx_t num_nonzeros = 0;
        // First Pass
        for (idx_t ic = ibeg; ic < iend; ic ++) {
            for (idx_t pa = A.row_ptr[ic]; pa < A.row_ptr[ic + 1]; pa ++) {
                idx_t j = A.col_idx[pa];
                marker[j] = ic;
                num_nonzeros ++;
            }
            for (idx_t pb = B.row_ptr[ic]; pb < B.row_ptr[ic + 1]; pb ++) {
                idx_t j = B.col_idx[pb];
                if (marker[j] != ic) {
                    marker[j]  = ic;
                    num_nonzeros ++;
                }
            }
            C.row_ptr[ic + 1] = num_nonzeros;
        }
        twspace[tid] = num_nonzeros;

        #pragma omp barrier
        // prefix sum
        if (tid > 0) {
            idx_t prev_sum = twspace[0];
            for (int pt = 1; pt < tid; pt ++) prev_sum += twspace[pt];
            for (idx_t ic = ibeg; ic < iend; ic ++) C.row_ptr[ic + 1] += prev_sum;
        } else {// 0号线程
            num_nonzeros = 0;// 记录所有非零元数总和
            for (int t = 0; t < nt; t ++) num_nonzeros += twspace[t];

            C.nnz = num_nonzeros;
            C.col_idx = new idx_t [C.nnz];
            C.vals    = new data_t[C.nnz * e_size];
        }

        #pragma omp barrier
        // Second Pass
        for (idx_t j = 0; j < C.ncols; j++) marker[j] = -1;// 初始化标记

        idx_t pos = C.row_ptr[ibeg];
        for (idx_t ic = ibeg; ic < iend; ic ++) {
            for (idx_t pa = A.row_ptr[ic]; pa < A.row_ptr[ic + 1]; pa ++) {
                idx_t j = A.col_idx[pa];
                C.col_idx[pos] = j;
                for (int f = 0; f < e_size; f++)
                    C.vals[pos * e_size + f] = alpha * A.vals[pa * e_size + f];
                marker[j] = pos;
                pos ++;
            }
            for (idx_t pb = B.row_ptr[ic]; pb < B.row_ptr[ic + 1]; pb ++) {
                idx_t j = B.col_idx[pb];
                if (marker[j] < C.row_ptr[ic]) {// 标记仍然是之前的某一行的
                    C.col_idx[pos] = j;
                    for (int f = 0; f < e_size; f++)
                        C.vals[pos * e_size + f] = beta * B.vals[pb * e_size + f];
                    marker[j] = pos;
                    pos ++;
                } else {// 标记是当前这一行的，说明此前j这一列已经有有效值，需要累加
                    assert(C.col_idx[marker[j]] == j);
                    for (int f = 0; f < e_size; f++)
                        C.vals[marker[j] * e_size + f] += beta * B.vals[pb * e_size + f];
                }
            }
            assert(pos == C.row_ptr[ic + 1]);
            sort_indices_with_vals<idx_t, data_t, dof>(C.row_ptr[ic + 1] - C.row_ptr[ic],
                C.col_idx + C.row_ptr[ic],
                C.vals    + C.row_ptr[ic] * e_size);// 结果矩阵的每一行按升序排列
        }
        assert(pos == C.row_ptr[iend]);

        delete [] marker;
    }// omp region
    delete [] twspace;
}

template<typename idx_t,    typename data_t, typename calc_t, int dof>
void CSRMatrixSplit(const CSRMatrix<idx_t, data_t, calc_t, dof> & B,
    const idx_t beg_col_B, const idx_t end_col_B, const idx_t offd_ncols_B, const idx_t * col_map_offd_B,
    idx_t & offd_ncols_C, idx_t * & col_map_offd_C,
    CSRMatrix<idx_t, data_t, calc_t, dof> & B_diag, CSRMatrix<idx_t, data_t, calc_t, dof> & B_offd)
{
    const idx_t nrows = B.nrows;
    constexpr int e_size = dof*dof;
    
    B_diag.nrows = nrows;
    B_diag.ncols = end_col_B - beg_col_B;
    B_diag.row_ptr = new idx_t [nrows + 1];

    B_offd.nrows = nrows;
    B_offd.row_ptr = new idx_t [nrows + 1];

    idx_t * my_diag_array = new idx_t [omp_get_max_threads()];
    idx_t * my_offd_array = new idx_t [omp_get_max_threads()];
    for (int t = 0; t < omp_get_max_threads(); t++) {
        my_diag_array[t] = my_offd_array[t] = 0;
    }

    idx_t * temp = nullptr;
    idx_t * B_offd_bigcid = nullptr;

    #pragma omp parallel
    {
        #pragma omp for schedule(static)
        for (idx_t i = 0; i < nrows + 1; i++) {
            B_diag.row_ptr[i] = 0;
            B_offd.row_ptr[i] = 0;
        }

        int tid = omp_get_thread_num();
        int  nt = omp_get_num_threads();
        idx_t ibeg = -1, iend = -1;
        thread_load_balance(nrows, nt, tid, ibeg, iend);

        idx_t my_diag_size = 0, my_offd_size = 0;
        for (idx_t i = ibeg; i < iend; i ++) {
            B_diag.row_ptr[i] = my_diag_size;
            B_offd.row_ptr[i] = my_offd_size;
            for (idx_t p = B.row_ptr[i]; p < B.row_ptr[i + 1]; p++) {
                if (beg_col_B <= B.col_idx[p] && B.col_idx[p] < end_col_B)// 位于对角部分
                    my_diag_size ++;
                else // 非对角部分
                    my_offd_size ++;
            }
        }
        my_diag_array[tid] = my_diag_size;
        my_offd_array[tid] = my_offd_size;
    
        #pragma omp barrier

        if (tid > 0) {
            my_diag_size = my_diag_array[0];
            my_offd_size = my_offd_array[0];
            for (int t = 1; t < tid; t ++) {
                my_diag_size += my_diag_array[t];
                my_offd_size += my_offd_array[t];
            }
            for (idx_t i = ibeg; i < iend; i++) {
                B_diag.row_ptr[i] += my_diag_size;
                B_offd.row_ptr[i] += my_offd_size;
            }
        } else {// 0号线程
            B_diag.nnz = 0;
            B_offd.nnz = 0;
            for (int t = 0; t < nt; t++) {
                B_diag.nnz += my_diag_array[t];
                B_offd.nnz += my_offd_array[t];
            }
            B_diag.row_ptr[nrows] = B_diag.nnz;
            B_offd.row_ptr[nrows] = B_offd.nnz;

            B_diag.col_idx = new idx_t [B_diag.nnz];
            B_diag.vals    = new data_t[B_diag.nnz * e_size];
            B_offd_bigcid  = new idx_t [B_offd.nnz];
            B_offd.col_idx = new idx_t [B_offd.nnz];
            B_offd.vals    = new data_t[B_offd.nnz * e_size];

            temp = new idx_t [B_offd.nnz + offd_ncols_B];// 按最大可能开辟
        }

        #pragma omp barrier

        idx_t cnt_diag = B_diag.row_ptr[ibeg];
        idx_t cnt_offd = B_offd.row_ptr[ibeg];
        for (idx_t i = ibeg; i < iend; i++) {
            for (idx_t p = B.row_ptr[i]; p < B.row_ptr[i + 1]; p++) {
                const idx_t j = B.col_idx[p];
                if (beg_col_B <= j && j < end_col_B) {// 对角部分
                    B_diag.col_idx[cnt_diag] = j - beg_col_B;// 减掉本进程偏移
                    for (int f = 0; f < e_size; f++)
                        B_diag.vals[cnt_diag * e_size + f] = B.vals[p * e_size + f];
                    cnt_diag ++;
                } else {// 非对角部分
                    temp[cnt_offd] = j;// 暂存全局序号
                    B_offd_bigcid[cnt_offd] = j;// 暂存全局序号
                    for (int f = 0; f < e_size; f++)
                        B_offd.vals[cnt_offd * e_size + f] = B.vals[p * e_size + f];
                    cnt_offd ++;
                }
            }
        }
        assert(cnt_diag == B_diag.row_ptr[iend]);
        assert(cnt_offd == B_offd.row_ptr[iend]);

        #pragma omp barrier
        // compute the mapping
        if (tid == 0) {
            idx_t cnt = B_offd.nnz;
            // 此时temp中 [0, B_offd.nnz) 的数据是有效的，后面的尚未填入
            // 因此将原来的 B矩阵的 col_map_offd 填入后面
            // 则 temp 中保存了原来的B和新split的B的所有全局列序号
            for (idx_t i = 0; i < offd_ncols_B; i ++) {
                temp[cnt ++] = col_map_offd_B[i];
            }
            assert(offd_ncols_C == 0);// 初始化要求
            if (cnt > 0) {
                my_qsort0(temp, 0, cnt - 1);// 快排升序排列（此时含有重复列序）
                offd_ncols_C = 1;// 记录去重后的有效的列数
                idx_t value = temp[0];
                for (idx_t i = 1; i < cnt; i++) {// 逐个考察可否往前递补
                    if (value < temp[i]) {// 是一个相异的值，可以往前递补
                        value = temp[i];
                        temp[offd_ncols_C ++] = value;
                    }
                }
            }

            col_map_offd_C = new idx_t [offd_ncols_C];
            for (idx_t i = 0; i < offd_ncols_C; i++)
                col_map_offd_C[i] = temp[i];

            delete [] temp; temp = nullptr;
        }

        #pragma omp barrier
        // 填入最终的非对角部分的局部列序
        for (idx_t i = ibeg; i < iend; i ++) {
            for (idx_t p = B_offd.row_ptr[i]; p < B_offd.row_ptr[i + 1]; p ++) {
                const idx_t glb_j = B_offd_bigcid[p];// 检索的目标值
                B_offd.col_idx[p] = my_BinarySearch(col_map_offd_C, glb_j, offd_ncols_C);
            }
        }
    }// omp region

    B_offd.ncols = offd_ncols_C;

    delete [] my_diag_array;
    delete [] my_offd_array;
    delete [] B_offd_bigcid;
}

template<typename idx_t,    typename data_t1, typename calc_t1,
                            typename data_t2, typename calc_t2, int dof>
void CSRMatrixCopy(const    CSRMatrix<idx_t, data_t1, calc_t1, dof> & src,
                            CSRMatrix<idx_t, data_t2, calc_t2, dof> & dst)
{
    dst.nrows = src.nrows;
    dst.ncols = src.ncols;
    dst.nnz   = src.nnz  ;
    dst.alloc_mem(false);
    constexpr int e_size = dof*dof;
    #pragma omp parallel
    {
        #pragma omp for schedule(static) nowait
        for (idx_t i = 0; i <= dst.nrows; i ++)
            dst.row_ptr[i] = src.row_ptr[i];
        
        #pragma omp for schedule(static) nowait
        for (idx_t k = 0; k <  dst.nnz  ; k ++) {
            dst.col_idx[k] = src.col_idx[k];
            for (int f = 0; f < e_size; f++)
                dst.vals[k * e_size + f] = src.vals[k * e_size + f];
        }
    }
}
#endif